import torch
from torchvision import transforms as transforms
import PIL

from mnist.mnist_training import alexnet

# 加载自己的样本数据
image = PIL.Image.open('./2.png')
image = image.resize((96,96))  # 输入尺寸与网络的输入保持一致
image = image.convert('L')     # 转为灰度图，保持通道数与网络输入一致
transform = transforms.Compose([transforms.Normalize((0.2, 0.2, 0.2,), (0.3, 0.3, 0.3,))]) # 对数据标准化的时候，标准化参数最好跟训练的数据一样
image = transform(image)

Infer_model = alexnet()  # 获得网络结构
Infer_model.load_state_dict(torch.load('alexnet_weight.pth')) # 加载最后训练的参数，在进行推理时，不需要优化器（optimizer），因为优化器只在训练时用于更新模型参数。

Infer_model.eval()  # 将模型转化为评估模型，此时虽然训练模型的batch_size是64，但是推理的时候可以只用一张图片
with torch.no_grad():
    output= Infer_model(image)  # 得到推理结果

# 返回函数的最大值的下标
pred = torch.argmax((output))
print('Prediction:', pred.item())
